#include "..\interface_types.h"
#include "d3d_depthmapshader.h"
#include "..\d3d_texture.h"
#include "..\d3d.h"
#include "..\d3d_mesh.h"
#include "..\d3d_util.h"
#include "..\..\common\string.h"
#include "..\..\common\def.h"
#include "..\..\engine\engine.h"
//#include "..\..\engine\memorymanager.h"
#include "..\..\engine\error.h"
#include "..\..\math\math.h"

#ifdef _DEBUG
#include "..\..\input\dinput.h"

#endif
D3DDepthMapShader::D3DDepthMapShader(CD3D9Device *pDevice) : D3DShader(pDevice)
{
	//m_Flags = 0; 
	mb_Locked = false; 
	mp_Technique = NULL; 
	mp_MaskTechnique = NULL; 

	m_lastInstanceCount = 0; 
	m_pInstanceVB = NULL; 

	mp_MaskInstanceVB = NULL; 
	mui_LastMaskInstanceCount = 0; 
}
string D3DDepthMapShader::GetType()
{	return "D3DDepthMapShader"; 
}

string D3DDepthMapShader::Identifier()
{
	return "D3DDepthMapShader"; 
}

void D3DDepthMapShader::PreExecute()
{
	mb_Locked = true; 

	if(FAILED(ConstructInstanceBuffer()))
		mb_Locked = false;
}

void D3DDepthMapShader::AddRenderQue(D3DRenderable* pMesh, const Matrix4f& transform)
{
#ifdef _DEBUG
	if(mb_Locked)
	{
		EE_LOGERROR("Attempted to add mesh while depth map shader was locked"); 
		return; 
	}
#endif

	if(!mb_Locked)
	{

		for(unsigned int numLights = GetDevice()->m_DirectionalLights.Count(), L = 0;  L < numLights;  L ++)
		{
			D3DDirectionalLight *pLight = GetDevice()->m_DirectionalLights[L];
			if(pLight)
			{
				const Matrix4f *pTransform = &transform;
				if(pLight->GetState() == D3DLight::LIGHT_ON && 
					ViewOrrientation::ContainsSphere(pLight->GetViewOrrientation(), 
					&Spheref(pTransform->_M41,pTransform->_M42,pTransform->_M43, pMesh->GetRadius())))
				{
					
					InstanceData instance; 
					instance.fData[0][0] = transform._M11; 
					instance.fData[0][1] = transform._M12; 
					instance.fData[0][2] = transform._M13; 
					instance.fData[0][3] = 0.0f; 

					instance.fData[1][0] = transform._M21; 
					instance.fData[1][1] = transform._M22; 
					instance.fData[1][2] = transform._M23; 
					instance.fData[1][3] = 0.0f; 

					instance.fData[2][0] = transform._M31; 
					instance.fData[2][1] = transform._M32; 
					instance.fData[2][2] = transform._M33; 
					instance.fData[2][3] = 0.0f; 

					instance.fData[3][0] = transform._M41; 
					instance.fData[3][1] = transform._M42; 
					instance.fData[3][2] = transform._M43; 
					instance.fData[3][3] = 1.0f; 

					if(pMesh->getMaskTexture())
						m_MaskMeshHashInstances.Add(pMesh, instance); 
					else
						m_MeshHashInstances.Add(pMesh, instance); 

					m_SubmittedVertexCount += pMesh->getVertexCount();
					return;
	
				}
			}
		}

	}
}

/**************************************
* Abstracts the task of instance buffer 
* creation
**************************************/
HRESULT D3DDepthMapShader::ConstructSubGroupInstanceBuffer(HashTable<InstanceData, D3DRenderable*>* pSubGroup, unsigned int* pLastInstanceCount, LPDIRECT3DVERTEXBUFFER9* pInstanceBuffer)
{
#ifdef _DEBUG
	if(ISNULL(pSubGroup) || ISNULL(pLastInstanceCount))
	{
		EE_LOGERROR("Bad parameters.");
		pSubGroup->Empty();
		*pInstanceBuffer = NULL;
		return E_FAIL;
	}
#endif

	LPDIRECT3DVERTEXBUFFER9 pIB = *pInstanceBuffer;
	unsigned int meshCount = 0; 
	unsigned int totalCount = 0; 
	pSubGroup->GetDataInfo(&meshCount, &totalCount); 

	if(meshCount == 0)
	{
		pSubGroup->ClearAllocatedSpace();
		*pInstanceBuffer = NULL;
		return S_OK;
	}

	D3DShader::InstanceData *pTmpInstTrans = NULL; 
	LPDIRECT3DDEVICE9 pDevice = GetDevice()->GetDevice(); 

	//Create instancing buffer for non-masked meshes
	if(totalCount > *pLastInstanceCount || !pIB)
	{
		*pLastInstanceCount = totalCount; 
		if(pIB) 
		{
			pIB->Release(); 
			pIB = NULL; 
		}
		if(pDevice->CreateVertexBuffer(totalCount * sizeof(D3DShader::InstanceData),0,0,
			D3DPOOL_DEFAULT,&pIB,NULL) != D3D_OK)
		{
			pSubGroup->Empty();
			*pInstanceBuffer = NULL;
			return E_FAIL;
		}
	}

	pTmpInstTrans = NULL; 
	if(SUCCEEDED(pIB->Lock(0, NULL, (void **) &pTmpInstTrans, 0)))
	{
		InstanceData* pMeshBaseByte = pTmpInstTrans; 
		for(unsigned int i = 0;  i < meshCount;  i++)
		{
			HashTable<InstanceData, D3DRenderable*>::Pair *pPair = pSubGroup->GetPair(i); 		
			memcpy(pMeshBaseByte, pPair->m_pData, pPair->m_Count * sizeof(D3DShader::InstanceData)); 
			pMeshBaseByte += pPair->m_Count; 
		}

		pIB->Unlock(); 
	}

	*pInstanceBuffer = pIB;
	return S_OK;
}
/**************************************
* Constructs an instance buffer for this
* shader based on the passed mesh instances
* in the last frame
**************************************/
HRESULT D3DDepthMapShader::ConstructInstanceBuffer()
{
	unsigned int vbSize = m_SubmittedVertexCount * VertexComplex::size;
	if(m_pMeshVB) 
	{
		m_pMeshVB->Release(); 
		m_pMeshVB = NULL; 
	}

	LPDIRECT3DDEVICE9 pDevice = GetDevice()->GetDevice();
	//Build up VB of all meshes (they share the same vertex declaration)
	if(pDevice->CreateVertexBuffer(vbSize,0,0, D3DPOOL_DEFAULT,&m_pMeshVB,NULL) != D3D_OK)
		return E_FAIL; 

	VertexComplex *pTmpMeshVertex = NULL; 
	if(m_pMeshVB->Lock(0, NULL, (void **) &pTmpMeshVertex, 0) != D3D_OK)
		return E_FAIL;

	unsigned int meshCount;
	unsigned int i;
	m_MeshHashInstances.GetDataInfo(&meshCount, NULL);
	for(i = 0;  i < meshCount;  i++)
	{
		D3DRenderable*	pMesh = m_MeshHashInstances.KeyAtTableIndex(i);
		unsigned int	meshVertexCount = pMesh->getVertexCount();
		memcpy(pTmpMeshVertex, pMesh->getVertices(), meshVertexCount * pMesh->m_ElementSize); 

		//TODO:fix bug with single triangles
		/*if(meshVertexCount < 9)
		{
		int k= 3;
		k++;
		}*/
		pTmpMeshVertex += meshVertexCount; 
	}
	m_MaskMeshHashInstances.GetDataInfo(&meshCount, NULL);
	for(i = 0;  i < meshCount;  i++)
	{
		D3DRenderable*	pMesh = m_MaskMeshHashInstances.KeyAtTableIndex(i);
		unsigned int	meshVertexCount = pMesh->getVertexCount();
		memcpy(pTmpMeshVertex, pMesh->getVertices(), meshVertexCount * pMesh->m_ElementSize); 
		pTmpMeshVertex += meshVertexCount; 
	}

	m_pMeshVB->Unlock(); 

	if(FAILED(ConstructSubGroupInstanceBuffer(&m_MeshHashInstances, &m_lastInstanceCount, &m_pInstanceVB)))
		return E_FAIL;
	if(FAILED(ConstructSubGroupInstanceBuffer(&m_MaskMeshHashInstances, &mui_LastMaskInstanceCount, &mp_MaskInstanceVB)))
		return E_FAIL;

	return S_OK;
}

void D3DDepthMapShader::RenderDepthMap(RenderTarget* pDepthBuffer, const ViewOrrientation* pView)
{
#ifdef _DEBUG
	if(!pDepthBuffer)
	{
		EE_LOGERROR("Bad parameters"); 
		return; 
	}
#endif
	PIX(L"Enter render depth map")

	if(mb_Locked)
	{
		LPDIRECT3DDEVICE9 pDevice = GetDevice()->m_pD3DDevice; 

		m_pEffect->SetMatrix("g_mDepthProjection",&pView->m_Projection); 
		m_pEffect->SetMatrix("g_mDepthView",&pView->m_View); 

		

		LPDIRECT3DSURFACE9 pOrigonalSurface = NULL; 
		pDevice->GetRenderTarget(0, &pOrigonalSurface); 

		pDevice->SetRenderState(D3DRS_ZENABLE, TRUE);
		pDevice->SetRenderState(D3DRS_ZWRITEENABLE, TRUE);
		pDevice->SetRenderState(D3DRS_ALPHABLENDENABLE, FALSE);
		pDevice->SetRenderState(D3DRS_CULLMODE, D3DCULL_CCW);

		pDevice->SetRenderTarget(0, pDepthBuffer->m_pSurface); 
		pDevice->Clear(0L, NULL, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER, D3DCOLOR_COLORVALUE(1.0f,1.0f,1.0f,1.0f), 1.0f, 0L); 

		pDevice->BeginScene(); 

		GetDevice()->SetCurrentVertexDeclaration(m_pVertexDeclaration);

		
		PIX(L"  Begin render opaqure meshes")
		pDevice->SetStreamSource(0, m_pMeshVB, 0, VertexComplex::size); 

		const unsigned int meshCount = m_MeshHashInstances.TableKeyCount(); 
		unsigned int uiInstanceBytesOffset = 0; 
		m_pEffect->SetTechnique(mp_Technique);  
		unsigned int vbBytesOffset = 0; 
		for(unsigned int i = 0;  i < meshCount;  i++)
		{
			HashTable<InstanceData, D3DRenderable*>::Pair *pPair = m_MeshHashInstances.GetPair(i); 
			D3DMesh* pMesh = (D3DMesh*)pPair->m_Key;
			GetDevice()->SetTexture(1, pMesh->getTexture(D3DShader::UNIT_B)); 
			pDevice->SetStreamSourceFreq(0, D3DSTREAMSOURCE_INDEXEDDATA | pPair->m_Count); 

			pDevice->SetStreamSource(1, m_pInstanceVB, uiInstanceBytesOffset, sizeof(D3DShader::InstanceData)); 
			pDevice->SetStreamSourceFreq(1, D3DSTREAMSOURCE_INSTANCEDATA | 1); 
			pDevice->SetIndices(pMesh->pIndexBuffer); 

			unsigned int passes = 0; 
			m_pEffect->Begin(&passes, 0); 
			for(unsigned int p = 0;  p < passes;  p++)
			{
				m_pEffect->BeginPass(p); 
				pMesh->Render(vbBytesOffset);
				m_pEffect->EndPass(); 
			}
			m_pEffect->End(); 

			uiInstanceBytesOffset += (pPair->m_Count * sizeof(D3DShader::InstanceData)); 
			vbBytesOffset += pMesh->getVertexCount();
		}

		PIX(L"  End opaque meshes, begin masked")
		//Renders all geometry that is subjected to masking
		const unsigned int maskedMeshCount = m_MaskMeshHashInstances.TableKeyCount(); 
		uiInstanceBytesOffset = 0; 
		m_pEffect->SetTechnique(mp_MaskTechnique);  
		for(unsigned int i = 0;  i < maskedMeshCount;  i++)
		{
			HashTable<InstanceData, D3DRenderable*>::Pair *pPair = m_MaskMeshHashInstances.GetPair(i); 
			D3DMesh* pMesh = (D3DMesh*)pPair->m_Key;  //TODO: not safe

			GetDevice()->SetTexture(0, pMesh->getTexture(D3DShader::UNIT_MASK)); 
			GetDevice()->SetTexture(1, pMesh->getTexture(D3DShader::UNIT_B)); 
			//pDevice->SetStreamSource(0, pMesh->pVertexBuffer, 0, pMesh->m_ElementSize); 
			pDevice->SetStreamSourceFreq(0, D3DSTREAMSOURCE_INDEXEDDATA | pPair->m_Count); 

			pDevice->SetStreamSource(1, mp_MaskInstanceVB, uiInstanceBytesOffset, sizeof(D3DShader::InstanceData)); 
			pDevice->SetStreamSourceFreq(1, D3DSTREAMSOURCE_INSTANCEDATA | 1); 
			pDevice->SetIndices(pMesh->pIndexBuffer); 

			unsigned int passes = 0; 
			m_pEffect->Begin(&passes, 0); 
			for(unsigned int p = 0;  p < passes;  p++)
			{
				m_pEffect->BeginPass(p); 
				pMesh->Render(vbBytesOffset);
				m_pEffect->EndPass(); 
			}
			m_pEffect->End(); 

			uiInstanceBytesOffset += (pPair->m_Count * sizeof(D3DShader::InstanceData)); 
			vbBytesOffset += pMesh->getVertexCount();
		}

		PIX(L"  End masked meshes, restore RT 0")
		pDevice->EndScene(); 
		pDevice->SetStreamSourceFreq(0, 1); 
		pDevice->SetStreamSourceFreq(1, 1); 

		
		pDevice->SetRenderTarget(0, pOrigonalSurface); 
	}

	PIX(L"End RenderDepthMap\n"); 
}


HRESULT D3DDepthMapShader::Reload()
{
	try {

		DWORD dwShaderFlags = D3DXFX_NOT_CLONEABLE; 
#ifdef _DEBUG
		dwShaderFlags |= D3DXSHADER_DEBUG; 
#endif
		if(m_Identity.length() == 0 || !m_pEffect)
			throw EE_EXCEPTION("")

			if(FAILED(m_pEffect->Release()))
				throw EE_EXCEPTION("")

				//Load the shader
				if(D3DXCreateEffectFromFile(GetDevice()->m_pD3DDevice,m_Identity.c_str(),
					NULL,NULL,dwShaderFlags,NULL,&m_pEffect,NULL) != D3D_OK)
					throw EE_EXCEPTION("")

//					m_pProjectionHandle = m_pEffect->GetParameterBySemantic( NULL, "PROJECTION" ); 
//		m_pViewHandle = m_pEffect->GetParameterBySemantic( NULL, "VIEW" ); 

		m_pTimeHandle = m_pEffect->GetParameterBySemantic( NULL, "TIME" ); 

		mp_Technique = m_pEffect->GetTechniqueByName("t_Shade"); 
		mp_MaskTechnique = m_pEffect->GetTechniqueByName("t_ShadeMask"); 

		return S_OK; 
	}catch(CError::CErrorData error)
	{
		CError::CError(error, ERRORREPORT_LOG);
		Release(); 
		return E_FAIL; 
	}
}

HRESULT D3DDepthMapShader::Create(const string& szFilename)
{
	LPD3DXBUFFER pErrors	= NULL; 
	DWORD dwShaderFlags = D3DXFX_NOT_CLONEABLE; 
	//string filepath = string(); 
	try {

		if(ISNULL(GetDevice()))
			throw EE_EXCEPTION("Shader has no graphics device."); 

		if(FAILED(Release()))
			throw EE_EXCEPTION("Failed to Release shader resources."); 

		if(szFilename.length() == 0)
			throw EE_EXCEPTION("Bad shader filename."); 

		if(ISNULL(m_pEffect)) {
			string filepath = CEngine::GetInstance()->FindFilepath(szFilename); 

			if(filepath.length() == 0)
				throw EE_EXCEPTION("Failed to locate file."); 

	#ifdef _DEBUG
			dwShaderFlags |= D3DXSHADER_DEBUG; 
	#endif
			//Load the shader
			if(D3DXCreateEffectFromFile(GetDevice()->m_pD3DDevice,filepath.c_str(),
				NULL,NULL,dwShaderFlags,D3DShader::GetEffectPool(),&m_pEffect,&pErrors) != D3D_OK)
			{
				char* pError =(char*)pErrors->GetBufferPointer(); 
				MessageBox(0,pError,0,0); 
				throw EE_EXCEPTION(pError); 
			}
		}

		//TODO: check that each is valid
		m_pProjectionHandle = m_pEffect->GetParameterBySemantic( NULL, "PROJECTION" ); 
		m_pViewHandle = m_pEffect->GetParameterBySemantic( NULL, "VIEW" ); 

		m_pTimeHandle = m_pEffect->GetParameterBySemantic( NULL, "TIME" ); 

		mp_Technique = m_pEffect->GetTechniqueByName("t_Shade"); 
		mp_MaskTechnique = m_pEffect->GetTechniqueByName("t_ShadeMask"); 

		m_Identity = szFilename; 

		if(FAILED(GetDevice()->GetDevice()->CreateVertexDeclaration(InstanceDecl, &m_pVertexDeclaration)))
			throw EE_EXCEPTION("CreateVertexDeclaration(..) failed."); 

		return S_OK; 

	}catch(CError::CErrorData errorData)
	{
		CError::CError(errorData); 
		Release(); 
		return E_FAIL; 
	}
}

void D3DDepthMapShader::RenderMeshes(const int flags)
{	}

void D3DDepthMapShader::PostExecute()
{	
	m_SubmittedVertexCount = 0;
	mb_Locked = false; 

	m_MaskMeshHashInstances.ClearAllocatedSpace(); 
	m_MeshHashInstances.ClearAllocatedSpace(); 
}

HRESULT D3DDepthMapShader::Release()
{
	m_IsValid = false; 
	//TODO: make sure inherited resources are relased
	m_MaskMeshHashInstances.Empty(); 
	mui_LastMaskInstanceCount = 0; 

	mp_Technique = NULL; 
	mp_MaskTechnique = NULL; 
	mb_Locked = false; 

	if(m_pEffect)
	{
		m_pEffect->Release(); 
		m_pEffect = NULL; 
	}

	if(!ISNULL(m_pVertexDeclaration))
	{
		m_pVertexDeclaration->Release();
		m_pVertexDeclaration = NULL;
	}

	if(!ISNULL(m_pInstanceVB))
	{
		m_pInstanceVB->Release();
		m_pInstanceVB = NULL;
	}
	if(!ISNULL(m_pMeshVB))
	{
		m_pMeshVB->Release();
		m_pMeshVB = NULL;
	}
	m_Identity.empty(); 

	return S_OK; 
}